import os, sys, re
sys.path.append(os.path.dirname(__file__))

from PIL import Image


# from pypdf import PdfReader, PdfWriter
from PyPDF2 import PdfReader, PdfWriter # 注意，和上面导入的是一样的
import fitz  # PyMuPdf
import pymupdf
from hos import *
from common import *
from natsort import natsorted, ns

class MyPdf:
    def __init__(self, path):
        if os.path.exists(path):
            self.path = path
            self.save_path = os.path.join(os.path.dirname(path), 'output.pdf')
    
    def marge(self, paths):
        '''
            合并paths指定的所有pdf文件
        '''
        writer = PdfWriter()
        save_path = os.path.join(os.path.dirname(paths[0]), 'output.pdf')
        for path in paths:
            reader = PdfReader(path)
            for index in range(len(reader.pages)):
                writer.add_page(reader.pages[index])
        with open(save_path, 'wb') as f:
            writer.write(f)
        return save_path

    def img2pdf(self, imgs):
        '''
            将文件夹内的所有图片 转换为一个pdf文件
        '''
        doc = pymupdf.open()
        for img in imgs:
            img = pymupdf.open(img)
            rect = img[0].rect # pic dimension
            pdfbytes = img.convert_to_pdf() # make a PDF stream
            img.close() # no longer needed
            imgPDF = pymupdf.open("pdf", pdfbytes) # open stream as PDF
            page = doc.new_page(width = rect.width, # new page with ...
                                height = rect.height) # pic dimension
            page.show_pdf_page(rect, imgPDF, 0) # image fills the page
        save_path = os.path.join(os.path.dirname(imgs[0]), 'output.pdf')
        doc.save(save_path)
        print("转换成功")

    def convert_to_black_white(self, path):
        '''
            灰度pdf
        '''
        # doc = fitz.open(path)
        doc = pymupdf.open(path)
        doc2 = pymupdf.open() # 输出文件
        for i, page in enumerate(doc):
            pix = page.get_pixmap(dpi=300)
            print(pix)
            mode = "RGBA" if pix.alpha else "RGB"
            img = Image.frombytes(mode, [pix.width, pix.height], pix.samples)
            print(img)
            img = img.convert("L")
            img.save(f"./pdf_output/{i}.png")

            img = pymupdf.open(f"./pdf_output/{i}.png") 
            rect = img[0].rect # pic dimension
            pdfbytes = img.convert_to_pdf() # make a PDF stream
            img.close() # no longer needed
            imgPDF = pymupdf.open("pdf", pdfbytes) # open stream as PDF
            page = doc2.new_page(width = rect.width, # new page with ...
                                height = rect.height) # pic dimension
            page.show_pdf_page(rect, imgPDF, 0) # image fills the page
            # psg.EasyProgressMeter("Import Images", i+1, imgcount) # show our progress
            os.remove(f"./pdf_output/{i}.png")
        save_path = OSrep_name(path, OSname(path)+"_黑白")
        doc2.save(save_path)
        return save_path
    
    def compress(self):
        doc = pymupdf.open(self.path)
        doc.save(self.save_path,
            garbage=4,  # 合并重复对象
             deflate=True,  # 压缩非流对象
             clean=True,  # 清理文档
             deflate_images=True,  # 压缩图像
             deflate_fonts=True)
        doc.close()
        print(f"文件已保存为{self.save_path}")



    def rotate(self, path, angle, start_page=None, end_page=None, page_list=[]):
        '''
            批量旋转pdf：
            1. 只指定start_page，则从start_page往后的页面全部旋转
            2. start，end都指定，则旋转范围内页面
            3. page_list，旋转列表中包含的页码
        '''
        output_path = OSrep_name(path, 'output')
        reader = PdfReader(path)
        writer = PdfWriter()
        for index in range(len(reader.pages)):
            if start_page and index+1 >= start_page:
                if end_page:
                    if index+1<=end_page:
                        writer.add_page(reader.pages[index].rotate(angle))
                else:
                    writer.add_page(reader.pages[index].rotate(angle))
            elif index+1 in page_list:
                writer.add_page(reader.pages[index].rotate(angle))
            else:
                writer.add_page(reader.pages[index])
        with open(output_path, 'wb') as f:
            writer.write(f)
        return output_path


    def replace_page(self, path, re_path, PageNum=None, RePageNum=None, angle=0):
        '''
            用re_path的第re_page_num页 替换 path中第page_num页
        '''
        reader = PdfReader(path) # 原文档
        re_reader = PdfReader(re_path) # 签字页
        writer = PdfWriter()
        if PageNum is None: # 默认替换最后一页
            PageNum = len(reader.pages)
        if RePageNum is None: # 默认使用签字文件的第一页
            RePageNum = 1

        if PageNum>0:
            PageNum -= 1
        if RePageNum>0:
            RePageNum-=1

        for index in range(len(reader.pages)):
            if index==PageNum:
                writer.add_page(re_reader.pages[RePageNum].rotate(angle)) # 默认不旋转签字页角度
            else:
                writer.add_page(reader.pages[index])
        output_dir = os.path.join(OS.dirname(path), 'sign_output')
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, os.path.basename(path))
        with open(output_path, 'wb') as f:
            writer.write(f)
        print(f"替换完成，文件保存在{output_path}")


def extract_sign_page(folder):
    '''
     从指定目录下，扣出pdf文件的最后一页（签字页）
    '''
    save_path = os.path.join(folder, 'sign.pdf')
    writer = PdfWriter()
    for cur, dirs, files in os.walk(folder):
        for file in files:
            if file.endswith(".pdf") and "3-Objective Evidence" not in cur:
                path = os.path.join(cur, file)      
                print(path)
                reader = PdfReader(path)
                index = -1 if "Verification Test Report" not in file else 3
                writer.add_page(reader.pages[index])
    with open(save_path, 'wb') as f:
        writer.write(f)
    return save_path
    

HELP = '''
    --merge: 合并指定文件夹下面所有的pdf文档, 指令："python pdf.py --merge xxx文件夹"
    --gray: 灰度化指定路径的pdf, 指令："python pdf.py --gray pdf文件路径"
    --rotate: 批量旋转指定页面, 指令："python pdf.py --rotate pdf文件路径"
    --img2pdf: 合并指定目录下的所有图片为一个pdf文件, 指令："python pdf.py --img2pdf 文件夹路径"
    --replace:  用签字文件的某一页替换原文件的某一页，签字页面可以指定旋转角度， 指令："python pdf.py --replace 原文件,新文件"

    注意： -1 为最后一页，1 为第一页
'''


def main():
    if '-h' in sys.argv or '--help' in sys.argv:
        print(HELP)
        return
    
    path = parse_argv('-f')
    pdf = MyPdf(path)

    merge_dir = parse_argv('--merge')
    # 合并指定文件夹下面所有的pdf文档
    if merge_dir:
        pdf.marge(
            natsorted([os.path.join(merge_dir, i) for i in os.listdir(merge_dir) if i.endswith('.pdf')], alg=ns.PATH)
        )

    # 灰度化指定路径的pdf
    gray_path = parse_argv('--gray')
    if gray_path:
        for path in argv_files(path=gray_path):
            pdf.convert_to_black_white(path)

    # 批量旋转指定页面
    rotate_path = parse_argv('--rotate')
    if rotate_path:
        start_page = input("请输入开始页码（或跳过）： ").strip('"')
        start_page = None if not start_page else int(start_page)
        end_page = input("请输入结束页码（或跳过）： ").strip('"')
        end_page = None if not end_page else int(end_page)
        page_list = input("请输入指定页码，用英文逗号隔开（或跳过）： ").strip('"')
        page_list = [int(i.strip()) for i in page_list.split(',')]
        angle = int(input("请输入旋转角度： ").strip('"'))
        pdf.rotate(rotate_path, angle, start_page, end_page, page_list) # 暂时从第0页开始，旋转90度

    # 合并指定目录下的所有图片为一个pdf文件
    imgs_path = parse_argv('--img2pdf')
    if imgs_path:
        imgs = natsorted([i for i in argv_files(imgs_path)], alg=ns.PATH)
        pdf.img2pdf(imgs)

    if '--replace' in sys.argv:
        # 用法: python pdf.py --replace
        # 说明: 用签字文件的某一页替换原文件的某一页，签字页面可以指定旋转角度
        input_path = input("请输入原始文件路径： ").strip().strip('"')
        sign_path = input("请输入签字文件路径： ").strip().strip('"')
        if not input_path or not sign_path:
            print("文件不存在")
        
        PageNum = int(input("请输入原始文件页码： ").strip('"'))
        RePageNum = int(input("请输入签字文件页码： ").strip('"'))
        angle = int(input("请输入签字文件旋转角度： ").strip('"'))
        print(f"用{OS.basename(sign_path)}的第{RePageNum}页替换{OS.basename(input_path)}的第{PageNum}页")
        pdf.replace_page(input_path, sign_path, PageNum=PageNum, RePageNum=RePageNum, angle=int(angle))
    
    # 压缩pdf
    if '--compress' in sys.argv:
        pdf.compress()
        
        


if __name__ == "__main__":
    main()


    
